package com.LeeCode;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * 在带权树网络中统计可连接服务器对数目
 */

public class Code3067 {
    public static void main(String[] args) {
        int[][] edges = {
                {0, 1, 1},
                {1, 2, 5},
                {2, 3, 13},
                {3, 4, 9},
                {4, 5, 2}
        };
        int signalSpeed = 1;
        System.out.println(Arrays.toString(new Code3067().countPairsOfConnectableServers(edges, signalSpeed)));
    }

    public int[] countPairsOfConnectableServers(int[][] edges, int signalSpeed) {
        int n = edges.length + 1;
        List<int[]>[] graph = new ArrayList[n];
        Arrays.setAll(graph, i -> new ArrayList<>());
        for (int[] edge : edges) {
            int x = edge[0], y = edge[1], weight = edge[2];
            graph[x].add(new int[]{y, weight});
            graph[y].add(new int[]{x, weight});
        }

        int[] ans = new int[n];
        for (int i = 0; i < n; i++) {
            if (graph[i].size() == 1) {
                ans[i] = 0;
                continue;
            }
            int sum = 0;
            for (int[] neighbor : graph[i]) {
                int cnt = dfs(neighbor[0], i, neighbor[1], graph, signalSpeed);
                ans[i] += cnt * sum;
                sum += cnt;
            }
        }
        return ans;
    }

    private int dfs(int cur, int from, int pathSum, List<int[]>[] graph, int signalSpeed) {
        int cnt = pathSum % signalSpeed == 0 ? 1 : 0;
        for (int[] neighbor : graph[cur]) {
            int to = neighbor[0];
            if (to == from) continue;
            cnt += dfs(to, cur, pathSum + neighbor[1], graph, signalSpeed);
        }
        return cnt;
    }
}
